import numpy as np
import pandas as pd
from sklearn import datasets
from sklearn.tree import DecisionTreeClassifier
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
from sklearn.tree import plot_tree
import matplotlib.pyplot as plt

# 设置matplotlib支持中文显示
plt.rcParams["font.family"] = ["SimHei", "WenQuanYi Micro Hei", "Heiti TC"]  # 尝试多种中文字体
plt.rcParams["axes.unicode_minus"] = False  # 正确显示负号

# 后续代码保持不变...

# -------------------------
#  1. 决策树：挑瓜数据集实验
# -------------------------
# （1）构建挑瓜数据集（示例结构，需根据实际数据补充）
# 特征：色泽、根蒂、敲声、纹理、脐部、触感、密度、含糖率
# 标签：好瓜（是=1，否=0）
melon_data = pd.DataFrame({
    '色泽': ['青绿', '乌黑', '浅白', '青绿', '乌黑', '浅白'],
    '根蒂': ['蜷缩', '蜷缩', '蜷缩', '稍蜷', '稍蜷', '稍蜷'],
    '敲声': ['浊响', '浊响', '沉闷', '浊响', '沉闷', '浊响'],
    '纹理': ['清晰', '清晰', '清晰', '稍糊', '稍糊', '稍糊'],
    '脐部': ['凹陷', '凹陷', '凹陷', '稍凹', '稍凹', '稍凹'],
    '触感': ['硬滑', '硬滑', '硬滑', '软粘', '软粘', '软粘'],
    '密度': [0.697, 0.774, 0.634, 0.608, 0.556, 0.403],
    '含糖率': [0.460, 0.376, 0.264, 0.318, 0.215, 0.237],
    '好瓜': [1, 1, 0, 0, 0, 1]
})

# （2）预处理：特征编码（文本特征转数值）
# 仅示例，实际需覆盖所有类别；也可使用 OneHotEncoder
for col in ['色泽', '根蒂', '敲声', '纹理', '脐部', '触感']:
    melon_data[col] = melon_data[col].astype('category').cat.codes

# （3）划分特征与标签
X_melon = melon_data.drop('好瓜', axis=1)
y_melon = melon_data['好瓜']

# （4）训练决策树
dt_melon = DecisionTreeClassifier(criterion='entropy', max_depth=3, random_state=42)
dt_melon.fit(X_melon, y_melon)

# （5）可视化决策树（可选）
plt.figure(figsize=(12, 8))
plot_tree(dt_melon, feature_names=X_melon.columns, class_names=['坏瓜', '好瓜'], filled=True)
plt.title("决策树（挑瓜数据集）")
plt.show()

# -------------------------
#  2. 随机森林：load_wine 数据集实验
# -------------------------
# （1）加载官方数据集
wine = datasets.load_wine()
X_wine, y_wine = wine.data, wine.target
X_train, X_test, y_train, y_test = train_test_split(X_wine, y_wine, test_size=0.2, random_state=42)


# （2）手动实现简易随机森林
class MyRandomForest:
    def __init__(self, n_estimators=10, max_depth=5, random_state=42):
        self.n_estimators = n_estimators  # 树的数量
        self.max_depth = max_depth  # 单棵树最大深度
        self.trees = []  # 保存所有决策树
        np.random.seed(random_state)  # 固定随机种子

    def _bootstrap_sample(self, X, y):
        """自助采样（Bootstrap）"""
        n_samples = X.shape[0]
        indices = np.random.choice(n_samples, n_samples, replace=True)
        return X[indices], y[indices]

    def fit(self, X, y):
        """训练随机森林：多棵决策树+Bootstrap+特征随机选择"""
        for _ in range(self.n_estimators):
            # 1. 自助采样
            X_boot, y_boot = self._bootstrap_sample(X, y)

            # 2. 随机选特征（简单实现：选 sqrt(特征数) 个）
            n_features = X_boot.shape[1]
            n_select = int(np.sqrt(n_features))
            feature_indices = np.random.choice(n_features, n_select, replace=False)
            X_boot = X_boot[:, feature_indices]

            # 3. 训练单棵决策树
            tree = DecisionTreeClassifier(max_depth=self.max_depth, random_state=42)
            tree.fit(X_boot, y_boot)
            self.trees.append((tree, feature_indices))  # 记录用了哪些特征

    def predict(self, X):
        """多棵树投票预测"""
        predictions = []
        for tree, feat_idx in self.trees:
            # 仅用训练时选中的特征
            X_sub = X[:, feat_idx]
            pred = tree.predict(X_sub)
            predictions.append(pred)

        # 投票：取多数类，这里改用numpy的方式实现
        pred_array = np.array(predictions).T
        result = []
        for row in pred_array:
            result.append(np.argmax(np.bincount(row)))
        return np.array(result)


# （3）训练并测试各模型
# --- 官方决策树 ---
dt_wine = DecisionTreeClassifier(random_state=42)
dt_wine.fit(X_train, y_train)
y_pred_dt = dt_wine.predict(X_test)

# --- 官方随机森林 ---
rf_sklearn = RandomForestClassifier(n_estimators=10, random_state=42)
rf_sklearn.fit(X_train, y_train)
y_pred_rf_sk = rf_sklearn.predict(X_test)

# --- 自定义随机森林 ---
my_rf = MyRandomForest(n_estimators=10, max_depth=5, random_state=42)
my_rf.fit(X_train, y_train)
y_pred_rf_my = my_rf.predict(X_test)  # 注意：需处理特征选择逻辑（如测试集也仅用对应特征）

# （4）评估准确率
print("===== 挑瓜数据集（决策树） =====")
print(f"训练集准确率: {dt_melon.score(X_melon, y_melon):.2f}")

print("\n===== Wine 数据集 =====")
print(f"决策树准确率: {accuracy_score(y_test, y_pred_dt):.2f}")
print(f"Sklearn 随机森林准确率: {accuracy_score(y_test, y_pred_rf_sk):.2f}")
print(f"自定义随机森林准确率: {accuracy_score(y_test, y_pred_rf_my):.2f}")